package com.yu.flow.demo;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;

public class SyncTaskExecutor implements TaskExecutor {
    @Override
    public TaskResult<?> execute(TaskConfig taskConfig, WorkflowContext context) {
        long startTime = System.currentTimeMillis();
        int retryCount = 0;
        int maxRetries = taskConfig.getRetries();

        while (retryCount <= maxRetries) {
            try {
                // 准备输入参数
                Map<String, Object> inputArgs = prepareInput(taskConfig, context);

                // 执行任务
                Object result;
                if (taskConfig.getTimeout() > 0) {
                    // 带超时的执行
                    result = executeWithTimeout(taskConfig, context, inputArgs);
                } else {
                    // 普通执行
                    result = taskConfig.getHandler().apply(context);
                }

                long executionTime = System.currentTimeMillis() - startTime;
                return new TaskResult<>(TaskStatus.SUCCESS, result, null, executionTime, retryCount);

            } catch (Exception e) {
                retryCount++;
                long executionTime = System.currentTimeMillis() - startTime;

                if (retryCount > maxRetries) {
                    System.err.println("Task " + taskConfig.getName() + " failed after " + retryCount + " retries: " + e.getMessage());
                    return new TaskResult<>(TaskStatus.FAILED, null, e, executionTime, retryCount - 1);
                }

                System.out.println("Task " + taskConfig.getName() + " failed, retrying (" + retryCount + "/" + maxRetries + "): " + e.getMessage());
                try {
                    Thread.sleep(taskConfig.getRetryDelay());
                } catch (InterruptedException ie) {
                    Thread.currentThread().interrupt();
                    return new TaskResult<>(TaskStatus.FAILED, null, ie, executionTime, retryCount - 1);
                }
            }
        }

        // 不会执行到这里，因为上面的循环会覆盖所有情况
        return new TaskResult<>(TaskStatus.FAILED, null, new RuntimeException("Unexpected execution path"), 0, 0);
    }

    private Map<String, Object> prepareInput(TaskConfig taskConfig, WorkflowContext context) {
        Map<String, Object> inputArgs = new HashMap<>();

        for (Map.Entry<String, String> entry : taskConfig.getInputMapping().entrySet()) {
            String paramName = entry.getKey();
            String dataKey = entry.getValue();
            Object value = context.getData(dataKey);

            if (value == null) {
                System.out.println("Missing data for key '" + dataKey + "' required by task '" + taskConfig.getName() + "'");
            }

            inputArgs.put(paramName, value);
        }

        return inputArgs;
    }

    private Object executeWithTimeout(TaskConfig taskConfig, WorkflowContext context, Map<String, Object> inputArgs) throws Exception {
        ExecutorService executor = Executors.newSingleThreadExecutor();
        Future<Object> future = executor.submit(() -> taskConfig.getHandler().apply(context));

        try {
            return future.get(taskConfig.getTimeout(), TimeUnit.MILLISECONDS);
        } catch (TimeoutException e) {
            future.cancel(true);
            throw new TimeoutException("Task execution timed out after " + taskConfig.getTimeout() + " ms");
        } finally {
            executor.shutdownNow();
        }
    }
}